
import sys
import indexing.featureindex
from summaryprocessor import SummaryProcessor, SinglePositionProcessor, SkipZeroes
from genome.genomereader import GenomeReader

QUALITY_ZERO = 64
# QUALITY_ZERO = 33

def incrFragCount(cts, (x, y)):
    cts[(x, y)] = cts.get((x, y), 0) + 1
    if x==None:
      # Look for (ANY, y)
      for (xx, yy) in cts.keys():
        if xx!=None and yy==y:
          cts[(xx, yy)] += cts[(x, y)]
          del cts[(x, y)]
          break
    if y==None:
      # Look for (x, ANY)
      for (xx, yy) in cts.keys():
        if xx==x and yy!=None:
          cts[(xx, yy)] += cts[(x, y)]
          del cts[(x, y)]
          break
    if y!=None and cts.has_key((x, None)):
      cts[(x, y)] += cts[(x, None)]
      del cts[(x, None)]
    if x!=None and cts.has_key((None, y)):
      cts[(x, y)] += cts[(None, y)]
      del cts[(None, y)]

class ReadsPerBase(SinglePositionProcessor):

  def __init__(self):
    SinglePositionProcessor.__init__(self)
    self.weightedCounts = 0.0
    self.fragments = {}
    self.lowQualityFragments = {}
    self.totalFragments = {}
    self.totalHighQualityFragments = {}
    self.totalReads = {}
    self.startRanges = {}  # range of start positions for (high-quality) covering reads showing base
    self.endRanges = {}    # range of end positions for (high-quality) covering reads showing base

    self.buffer = []
    self.genomicBase = None
    self.hqNonConsensReads = 0
    
  def incr(self, basequal, seqpos, details):
    (base, qual) = basequal
    qual = ord(qual)-QUALITY_ZERO
    if qual>40 or qual<-5:
        raise Exception, "quality value %d out of range (bad QUALITY_ZERO=%d ?)" % (qual, QUALITY_ZERO)
    if self.hqNonConsensReads < 2:
    # if self.hqNonConsensReads < 0:
        if base != self.genomicBase and qual >= 20:
            self.hqNonConsensReads += 1
        if self.hqNonConsensReads >= 2:
            for (a1, a2, a3) in self.buffer:
                self.incr(a1, a2, a3)
        else:
            self.buffer.append((basequal, seqpos, details))
            return

    SinglePositionProcessor.incr(self, base, seqpos, details)
    
    if qual>=20:
      (minpos, maxpos) = details.getReadEndpoints()
      if not self.startRanges.has_key(base):
        self.startRanges[base] = (minpos, minpos)
      else:
        rng = self.startRanges[base]
        self.startRanges[base] = (min(rng[0], minpos), max(rng[1], minpos))
      if not self.endRanges.has_key(base):
        self.endRanges[base] = (maxpos, maxpos)
      else:
        rng = self.endRanges[base]
        self.endRanges[base] = (min(rng[0], maxpos), max(rng[1], maxpos))
    
    if not details.spliceInterval:
      seqlen = len(details.seq)
    else:
      seqlen = len(details.seq) + int(details.startDrop) + int(details.endDrop)

    weight = 1.0 / seqlen
    self.weightedCounts += weight

    if qual>=20:
      frags = self.fragments
    else:
      frags = self.lowQualityFragments

    if not frags.has_key(base):
      frags[base] = {}

    # this is the unique key for fragments... if two reads have the
    # same value for this key, they are counted as the same fragment

    # key = details.getLaneId() + "/" + details.getFragmentSpan()
    (X, Y) = details.getFragmentEndpoints()
    if details.libraryId:
      if X:
        X = "%d(%s)" % (X, details.libraryId)
      if Y:
        Y = "%d(%s)" % (Y, details.libraryId)
    
    # frags[base][key] = frags[base].get(key, 0) + 1
    incrFragCount(frags[base], (X, Y))

    # self.totalFragments[key] = self.totalFragments.get(key, 0) + 1
    incrFragCount(self.totalFragments, (X, Y))

    self.totalReads[base] = self.totalReads.get(base, 0) + 1
    if qual >= 20:
      # self.totalHighQualityFragments[key] = self.totalHighQualityFragments.get(key, 0) + 1
      incrFragCount(self.totalHighQualityFragments, (X, Y))
  
  def hasDiscrepancy(self, base):
    return self.totalReads.get(base, 0) < self.total

  def getFragments(self, filtered=True):
    if filtered:
      return len(self.totalHighQualityFragments)
    else:
      return len(self.totalFragments)
  
  def getReads(self, filtered=True):
    if filtered:
      return sum([val for (frag, val) in self.totalHighQualityFragments.items()])
    else:
      return sum([val for (frag, val) in self.totalFragments.items()])

  def getFragmentCounts(self, filtered=True):
    cts = {'A':0, 'C':0, 'G':0, 'T':0, 'N':0}
    for (base, frags) in self.fragments.items():
      cts[base] = len(frags)
    if not filtered:
      for (base, frags) in self.lowQualityFragments.items():
        cts[base] += len(frags)
    return tuple([cts[base] for base in ['A','C','G','T']])
  
  def getReadCounts(self, filtered=True):
    cts = {'A':0, 'C':0, 'G':0, 'T':0, 'N':0}
    for (base, frags) in self.fragments.items():
      cts[base] = sum([val for (frag, val) in frags.items()])
    if not filtered:
      for (base, frags) in self.lowQualityFragments.items():
        cts[base] += sum([val for (frag, val) in frags.items()])
    return tuple([cts[base] for base in ['A','C','G','T']])

  def getPositionRanges(self):
    sRanges = {'A':0, 'C':0, 'G':0, 'T':0}
    eRanges = {'A':0, 'C':0, 'G':0, 'T':0}
    for base in self.startRanges.keys():
      sRanges[base] = self.startRanges[base][1] - self.startRanges[base][0]
      eRanges[base] = self.endRanges[base][1] - self.endRanges[base][0]
    return "%s/%s/%s/%s\t%s/%s/%s/%s" % tuple([sRanges[base] for base in 'ACGT'] + [eRanges[base] for base in 'ACGT'])

  def __str__(self):
    return "%d\t%.6f\t%s" % (self.total, self.weightedCounts, self.fragments)

class Counts(SinglePositionProcessor):
  def __str__(self):
    return str(self.total)

def process(query, indexFilename, skipZeroes=(SkipZeroes.LEFT | SkipZeroes.RIGHT), positionChecker=None):
  processor = SummaryProcessor(ReadsPerBase, Counts, positionChecker)
  return processor.process(query, indexFilename, skipZeroes)

def printAsSnps(query, indexFilename, minVariantFrags=2, positionChecker=None):
  (qchr, qstrand, qmin, qmax) = query
  # reader = GenomeReader(qchr)
  for (x, cts) in process(query, indexFilename, SkipZeroes.ALL, positionChecker):
    # genomic = reader[x]
    if isinstance(cts, Counts): continue
    genomic = cts.genomicBase
    fragCounts = cts.getFragmentCounts()
    if not cts.hasDiscrepancy(genomic): continue
    fragCounts = cts.getFragmentCounts()
    readCounts = cts.getReadCounts()
    disc = "T"
    fragString = "%d/%d/%d/%d" % fragCounts
    readString = "%d/%d/%d/%d" % readCounts
    variantFrags = sum(fragCounts) - fragCounts['ACGT'.index(genomic)]
    hqSnp = (variantFrags > 0)
    if variantFrags < minVariantFrags: continue
    allFragCounts = cts.getFragmentCounts(False)
    allReadCounts = cts.getReadCounts(False)
    allFragString = "%d/%d/%d/%d" % allFragCounts
    allReadString = "%d/%d/%d/%d" % allReadCounts
    anyVariantReads = sum(allReadCounts) - allReadCounts['ACGT'.index(genomic)]
    if not anyVariantReads:
        disc = "F"
    print "%s\t%ld\t%s\t%s\t%ld\t%.6f\t%ld\t%s\t%s\t%s\t%s\t%s\t%s" % (qchr, x, hqSnp, genomic, cts.total, cts.weightedCounts, len(cts.totalFragments), disc, readString, fragString, cts.getPositionRanges(), allReadString, allFragString)

def printAsCoverageSummary(query, indexFilename):
  printedInitialStep = False
  (qchr, qstrand, qmin, qmax) = query
  reader = GenomeReader(qchr)
  prevPos = 0
  for (x, cts) in process(query, indexFilename, SkipZeroes.ALL):
    genomic = reader[x]
    if x != prevPos+1 or not printedInitialStep:
      print "fixedStep chrom=%s start=%ld step=1" % (qchr, x)
      printedInitialStep = True
    prevPos = x
      
    fragCounts = cts.getFragmentCounts()
    readCounts = cts.getReadCounts()
    if cts.hasDiscrepancy(genomic):
      disc = "T"
      fragString = "%d/%d/%d/%d" % fragCounts
      readString = "%d/%d/%d/%d" % readCounts
    else:
      disc = "F"
      fragString = str(sum(fragCounts))
      readString = str(sum(readCounts))
    print "%s\t%ld\t%.6f\t%ld\t%s\t%s\t%s" % (genomic, cts.total, cts.weightedCounts, len(cts.totalFragments), disc, readString, fragString)

def printUsage(args, usage, minArgs=2):
  if len(args) < minArgs or args[1].startswith("-h") or args[1].startswith("--h"):
    print usage
    sys.exit(1)

usageString = """
Usage: python %s INDEX_FILE CHR MIN_POS MAX_POS

Prints the sum of mapping weights in the specified summary file at each position
in the query.
Mapping weights are 1/N for each mapping, where N is the number of locations
to which the read was mapped.
"""

chrs = ["chr"+str(x) for x in xrange(1,23)] + ["chrX", "chrY"]

if __name__=="__main__":
  printUsage(sys.argv, usageString % sys.argv[0], 3)
  asSnps = False
  if sys.argv[1]=="--as-snps":
    asSnps = True
    sys.argv = sys.argv[1:]
  if sys.argv[1]=="-ALL":
    queries = ((chr, '?', 0, 1000000000L) for chr in chrs)
    indexFilename = sys.argv[2]
  else:
    indexFilename = sys.argv[1]
    queries = []
    if sys.argv[2] == '-F':
        qstrs = []
        for line in open(sys.argv[3]):
            qstrs.append(line.strip())
        print >> sys.stderr, "# Read %d queries from file %s" % (len(qstrs), sys.argv[3])
    else:
        qstrs = sys.argv[2:]
        print >> sys.stderr, "# Read %d queries from command line" % len(qstrs)
    for qstr in qstrs:
        if not ":" in qstr:
            queries.append((qstr, '?', 0, 1000000000L))
        else:
            (chr, ival) = qstr.split(":")
            (minpos, maxpos) = ival.split("-")
            minpos = long(minpos)
            maxpos = long(maxpos)
            queries.append((chr, '?', minpos, maxpos))
    print >> sys.stderr, "# Processing", len(queries), "queries: ", queries[:10]
  
  indexFilename = indexing.featureindex.FeatureIndex(indexFilename)
  
  nq = 0
  for query in queries:
    if asSnps:
      printAsSnps(query, indexFilename)
    else:
      printAsCoverageSummary(query, indexFilename)
      nq += 1
      print >> sys.stderr, "# Query", nq, ":", query
